from __future__ import print_function
from torch.utils.data import Dataset
from torch.autograd import Variable
import torch
import numpy as np
import pickle


def MdataLoader(path):
    # load data
    # data shape = [samples, max_length]
    f = open(path, 'rb')
    data = pickle.load(f)
    f.close()
    return data


class Mdataset(Dataset):
    # dataset class for text classification
    def __init__(self, path):
        self.data = MdataLoader(path)
        self.src_ = self.data["data"]
        self.target_ = self.data["label"]
        self.src_len_ = self.data["src_len"]
        assert len(self.src_) == len(self.target_)
        assert len(self.src_) == len(self.src_len_)
        # print(len(self.src_))

    def __len__(self):
        return len(self.src_)

    def __getitem__(self, indx):
        # getitem method
        # print("label ", self.target_[indx])
        self.sample_ = Variable(torch.LongTensor(self.src_[indx, ...])), \
                            Variable(torch.LongTensor(self.target_[indx])),\
                            Variable(torch.LongTensor(self.src_len_[indx]))

        return self.sample_

